'''
图片预处理脚本
1. resize224x224
2. 生成annotations.json
    bbox 处理： 对于点状标注，计算其中心点和半径，转换为bbox格式
zym1105 25/7/18
'''
import json
import os
from PIL import Image
from tqdm import tqdm
import numpy as np

data_root = '/home/guoyi/Dataset/aier_orignal/'  # 数据根目录
json_path = os.path.join(data_root, 'final_sign_result.json')
tar_dir = '/home/guoyi/Dataset/aier_processed/'    # 结果目录

os.makedirs(tar_dir, exist_ok=True)
os.system(f'rm -rf {tar_dir}/*')  # 清空结果目录
os.makedirs(os.path.join(tar_dir, 'images'), exist_ok=True)

with open(json_path, 'r', encoding='utf-8') as f:
    json_data_ori = json.load(f)
json_data = {}
for key, value in json_data_ori.items():
    # 把key的' '替换为'_'
    new_key = key.replace(' ', '_')
    json_data[new_key] = value
image_label_folder = os.listdir(data_root)
image_label_folder = [folder for folder in image_label_folder if os.path.isdir(os.path.join(data_root, folder))]

annotations = {}

# 统计信息
class_count = {}
lesion_count_per_image = []

for folder in image_label_folder:
    image_names = os.listdir(os.path.join(data_root, folder))
    class_count[folder] = len(image_names)
    for image_name in tqdm(image_names, desc=f'Processing {folder}'):
        image_path = os.path.join(data_root, folder, image_name)
        # 判断图片是否真的存在于json_data中（json_data键可能是绝对路径或相对路径，建议统一一下路径）
        if image_name in json_data:
            
            # 打开图片看大小
            image = Image.open(image_path)
            width, height = image.size
            # 保存图片到目标目录，并调整大小
            resized_image = image.resize((224, 224))
            resized_image.save(os.path.join(tar_dir, 'images', image_name))
            image.close()

            # 处理lesion
            lesions = []
            for item in json_data[image_name]:
                try:
                    lesion_name = item['props']['name']
                    shape = item['shape']
                    if 'width' in shape and 'height' in shape:
                        # 普通矩形框
                        x = shape['x'] / width
                        y = shape['y'] / height
                        w = shape['width'] / width
                        h = shape['height'] / height
                        bbox = [x, y, x + w, y + h]
                    elif 'sr' in shape:  # 处理点状
                        # 以点为中心，sr为半径
                        x = (shape['x'] - shape['sr']) / width
                        y = (shape['y'] - shape['sr']) / height
                        w = (2 * shape['sr']) / width
                        h = (2 * shape['sr']) / height
                        bbox = [x, y, x + w, y + h]
                    else:
                        print(f"Unknown lesion shape in {image_name}: {json.dumps(item, ensure_ascii=False, indent=2)}")
                        continue
                    lesions.append({
                        "name": lesion_name,
                        "bbox": bbox
                    })
                except KeyError as e:
                    print(f"KeyError in {image_name}: {e}. Skipping this lesion item.")
                    print(f"Item data: {json.dumps(item, ensure_ascii=False, indent=2)}")
                    raise ValueError(f"Invalid lesion data in {image_name}: {item}")
            annotations[image_name] = {
                "image_path"    : os.path.join(tar_dir, 'images', image_name),
                "image_rela_dir": os.path.join(folder, image_name),
                "eye_type": "left" if 'l' in image_name else "right",
                "lesions": lesions,
                "image_label": folder
            }
            lesion_count_per_image.append(len(lesions))
        else:
            print(f"Warning: {image_path} not found in json data.")

# 写入新json
with open(os.path.join(tar_dir, 'annotations.json'), 'w', encoding='utf-8') as f:
    json.dump(annotations, f, indent=4, ensure_ascii=False)

# 统计信息
print("\n=== 统计信息 ===")
print("各类别图片数：")
for k, v in class_count.items():
    print(f"{k}: {v}")

if lesion_count_per_image:
    print(f"\n每张图片的lesion数量:")
    print(f"平均值: {np.mean(lesion_count_per_image):.2f}")
    print(f"中位数: {np.median(lesion_count_per_image):.2f}")
    print(f"标准差: {np.std(lesion_count_per_image):.2f}")
else:
    print("没有有效的lesion统计数据。")
